import numpy as np
from PIL import Image
import matplotlib.pyplot  as plt
import matplotlib
from scipy import signal
from cv2 import GaussianBlur
import time
matplotlib.rcParams['font.family'] = 'FangSong'
matplotlib.rcParams['font.size'] = 7


# 高斯模糊的矩阵 3*3
def MyGaussianBlur(sigma):
    sigma = float(sigma)
    num1 = np.around((2 * np.pi * sigma ** 2) ** (-1), decimals=7)
    num2 = np.around((2 * np.pi * sigma ** 2) ** (-1) * np.exp((np.negative(sigma ** 2)) ** (-1) * 0.5), decimals=7)
    num3 = np.around((2 * np.pi * sigma ** 2) ** (-1) * np.exp((np.negative(sigma ** 2)) ** (-1)), decimals=7)
    GaussMatrix = np.array([[num3, num2, num3],
                            [num2, num1, num2],
                            [num3, num2, num3]])
    k = np.around(((num2 + num3) * 4 + num1), decimals=7)
    return GaussMatrix/k

# 返回浮点数

def Gauss_calculate(img,sigma):
    Gauss = np.zeros(shape=img.shape,dtype=float)
    M = MyGaussianBlur(sigma)
    for i in range(3):
        Gauss[..., i] = signal.convolve2d(img[..., i], M, mode="same", fillvalue=0)
    return Gauss

# 返回的是整数
def Gauss_calculate(img,sigma,ksize):
    Gauss = GaussianBlur(img, (ksize, ksize), sigma)
    return Gauss
# 因为会有取对数的过程所以要避免，图像数据中有0的像素点
def nozero(img):
    img[img == 0] = 1
    return img

def SSR(img,c):
    R = np.zeros(shape=img.shape,dtype=int)
    Gauss = Gauss_calculate(img,c/np.sqrt(2),3)
    img = nozero(img)
    Gauss = nozero(Gauss)
    r = np.log(img) + np.log(Gauss)
    for i in range(3):
        MAX = np.max(r[:,:,i])
        MIN = np.min(r[:,:,i])
        R[:,:,i] = (r[:,:,i] - MIN) / (MAX - MIN) * 255
    return R

def MSR(img,C,w = [1/3,1/3,1/3]):
    RR = np.zeros(shape=img.shape, dtype=float)
    img = nozero(img)
    for i in range(len(C)):
        Gauss = Gauss_calculate(img, C[i] / np.sqrt(2), 3)
        R = np.zeros(shape=img.shape, dtype=float)
        Gauss = nozero(Gauss)
        r = np.log(img) + np.log(Gauss)
        for i in range(3):
            MAX = np.max(r[:, :, i])
            MIN = np.min(r[:, :, i])
            R[:, :, i] = (r[:, :, i] - MIN) / (MAX - MIN) * 255
        RR += w[i]*R
    RR = RR.astype(int)
    return RR
def MSRCR(img,C,w =[1/3,1/3,1/3],Dynamic = 2):
    RR = np.zeros(shape=img.shape, dtype=float)
    img = nozero(img)
    for i in range(len(C)):
        Gauss = Gauss_calculate(img, C[i] / np.sqrt(2), 5)
        R = np.zeros(shape=img.shape, dtype=float)
        Gauss = nozero(Gauss)
        r = np.log(img) + np.log(Gauss)
        r = r.astype(np.float64)
        for i in range(3):
            MEAN = np.mean(r[:,:,i])
            VAR = np.std(r[:,:,i])
            MAX = MEAN + Dynamic * VAR
            MIN = MEAN - Dynamic * VAR
            R[:, :, i] = (r[:, :, i] - MIN) / (MAX - MIN) * 255
            R[R[:,:,i] > 255,i] = 255
            R[R[:, :, i] < 0, i] = 0
        RR += w[i] * R
    RR = RR.astype(int)
    return RR

#E:\pythonPro\Underwater\underwater\data\image\2019v100010.jpg

'''
SSR c 建议 80 - 100
尺度取值较小时,能够较好地完成动态范围的压缩，暗区域的细节能得到较好地增强，但输出颜色易失真；取值较大时，色感一致性较好
MSR c 建议 15,80,120
MSRCR
Dynamic取值越小，图像的对比度越强。
一般来说Dynamic取值2-3之间能取得较为明显的增强效果，即能取得很自然过渡效果，又能保持图像的清晰度适度增强。
关于最大尺度，建议取值以大于100为佳。
'''
t1 = time.time()

img = Image.open(r"E:\pythonPro\Underwater\underwater\data\image\2019v100179.jpg")
img_data = np.array(img)

para = (20,80,200)

new_img1 = MSRCR(img_data,para[:3],Dynamic=3)
#new_img2 = MSRCR(img_data,para[:3],Dynamic=4)
#new_img3 = MSRCR(img_data,para[:3],Dynamic=5)


t2 = time.time()
print('用时：',(t2-t1))


f = plt.figure()
plt.subplot(2,2,1)
plt.title("原图")
plt.imshow(img_data)

f.add_subplot(2, 2, 2)
plt.imshow(new_img1)
plt.title("MSRCR"+"c1:%d,c2:%d,c3:%d, Dynamic:3"%para)
plt.show()

img = Image.fromarray(np.uint8(new_img1))

img.save('test2.jpg')
'''
f.add_subplot(2, 2, 3)
plt.imshow(new_img2)
plt.title("MSRCR"+"c1:%d,c2:%d,c3:%d, Dynamic:4"%para)

f.add_subplot(2, 2, 4)
plt.imshow(new_img3)
plt.title("MSRCR"+"c1:%d,c2:%d,c3:%d, Dynamic:5"%para)

plt.show()
'''
